package com.example.springai.controller;

import com.example.springai.common.R;
import com.example.springai.domain.dto.VectorSearchDTO;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.math.BigDecimal;
import java.util.List;
import java.util.Objects;

/**
 * 向量测试控制器
 *
 * @author ZWX
 */
@RestController
@RequestMapping("/vector/test")
@RequiredArgsConstructor
public class VectorTestController {
    private final VectorStore vectorStore;

    /**
     * 搜索
     *
     * @return 搜索结果
     */
    @GetMapping("/search")
    public R<List<Document>> search(VectorSearchDTO req) {
        List<Document> documents = vectorStore.similaritySearch(req.getQuery());
        // 筛选分数
        Double scoreThreshold = req.getScoreThreshold();
        if (Objects.nonNull(scoreThreshold)) {
            documents = documents.stream().filter(document -> document.getScore().compareTo(scoreThreshold) > 0).toList();
        }
        // 筛选数量
        Integer topK = req.getTopK();
        if (Objects.nonNull(topK)) {
            documents = documents.stream().limit(topK).toList();
        }
        return R.ok(documents);
    }

}
